var  $ = require("jquery");
var d3 = require("d3");
var size = 400; //The size of the canvas
var margin_size = 50;
var point_size = 8;
var colors = d3.scaleOrdinal(d3.schemeCategory10);
var domain_max = 100;
var data = [];
var regression_model = undefined;
import * as tf from '@tensorflow/tfjs';

function drawCircle(container, p, r, color) {
  var circle = container
    .append("circle")
    .attr("cx", p.x)
    .attr("cy", p.y)
    .attr("r", r)
    .classed("circle", true);

  if (color) {
    circle.style("fill", color);
  }
  return circle;
}

function drawLine(container, p1, p2, color, delay) {
  var line = container
    .append("line")
    .classed("line", true)
    .attr("x1", p1.x)
    .attr("y1", p1.y)
    .attr("x2", p2.x)
    .attr("y2", p2.y)
    .style("stroke","none")
    .style("stroke-opacity",0.1)
    .transition()
    .delay(delay)
    .style("stroke","#ccc");
    
  return line;
}

function clean() {
  data = [];
  $(".circle").remove();
  $(".line").remove();
}

function linear_regression(tx, ty) {
  var results = [];
  const w = tf.variable(tf.scalar(Math.random()));
  const b = tf.variable(tf.scalar(Math.random())); 
  
  const numIterations = 200;
  const learningRate = 1;
  const optimizer = tf.train.adam(learningRate);
  
  //modle 
  const f = x => w.mul(x).add(b);

  //lost func
  //sub 减法 (（pred - label）^2)平均值
  const loss = (pred, label) => pred.sub(label).square().mean();
  
  const train_x = tf.tensor1d(tx);
  const train_y = tf.tensor1d(ty);
  
  for (let iter = 0; iter < numIterations; iter++) {
    optimizer.minimize(() => {
      const loss_var = loss(f(train_x), train_y);
      loss_var.print();
      return loss_var;
    })
    let result = {};
    result.w = w.dataSync()[0];
    result.b = b.dataSync()[0];
    result.f = x => result.w*x + result.b ; 
    results.push(result);
  }
  return results;
}

$(function() {
  var margin = {
    top: margin_size,
    right: margin_size,
    bottom: margin_size,
    left: margin_size
  },
    width = size - margin.left - margin.right,
    height = size - margin.top - margin.bottom;

  var root = d3
    .select("#chart")
    .append("svg")
    .attr("width", size)
    .attr("height", size);

  var g = root
    .append("g")
    .attr("transform", "translate(" + margin.left + "," + margin.top + ")");

  //Draw Axis
  var xScale = d3.scaleLinear().rangeRound([0, width]);
  var yScale = d3.scaleLinear().rangeRound([height, 0]);

  xScale.domain([0, domain_max]);
  yScale.domain([0, domain_max]);

  g
    .append("g")
    .attr("transform", "translate(0," + height + ")")
    .call(d3.axisBottom(xScale));

  g.append("g").call(d3.axisLeft(yScale));

  // Generate Data 获取数据
  $("#new_button").click(function() {
    clean();
    root.on("click", function() {
      var coords = d3.mouse(this);
      var mapped_coords = [coords[0] - margin_size, coords[1] - margin_size];
      var newData = {
        x: Math.round(xScale.invert(mapped_coords[0])), // Takes the pixel number to convert to number
        y: Math.round(yScale.invert(mapped_coords[1]))
      };
      drawCircle(g, { x: mapped_coords[0], y: mapped_coords[1] }, point_size);
      data.push(newData);
    });
  });

  //线性回归 train 开始
  $("#regression_button").click(function() {
    var data_x = [];
    var data_y = []
    //Turn data into array
    data.map(function(d) {
      data_x.push(d.x);
      data_y.push(d.y);
    });

    regression_model = linear_regression(data_x,data_y);
    let final_result = regression_model[regression_model.length-1];
    
    var delay = 0;
    var delay_diff = 10;
    var final_line = undefined;
    regression_model.map(function(model){
      var p1 = { x: 0, y: yScale(model.b) };
      var p2 = { x: xScale(domain_max), y: yScale(model.f(domain_max)) };
      final_line = drawLine(g, p1, p2, "#ccc", delay);
      console.log("delay " + delay);
      delay = delay + delay_diff;
    })
    
    final_line.transition()
    .delay(delay)
    .style("stroke","#0366d6")
    .style("stroke-opacity",0.8);
    
  });
  
  //clean
  $("#clean_button").click(function() {
    clean();
  });
});
